ICLR 2021｜一种端到端的基于双重优化的分子构象生成框架ConfVAE

本文由加拿大魁北克人工智能研究所唐建、Minkai Xu、Yoshua Bengio 以及 MIT 的 Wujie Wang、北京大学的Shitong Luo 等人合作完成，发表在2021年 ICLR（人工智能领域顶级会议）。

摘要

从分子图预测分子构象（或三维结构）是许多应用中的一个基本问题。大多数现有的方法通常分为两个步骤，首先预测原子之间的距离，然后通过优化距离几何问题生成三维结构。然而，用这种两阶段方法预测的构象可能无法保持真实原子的邻接关系，使得生成的结构不令人满意。在本文中，我们提出了一种基于条件变分自编码器 ConfVAE 的端到端分子构象预测模型。具体来说，首先将分子图编码到一个隐空间中，然后通过求解一个双层优化程序生成三维结构。在几个基准数据集上进行的大量实验证明了我们提出的方法比现有的最先进的方法更有效。

1 简介

人工智能交叉分子建模方兴未艾（从分子性质预测到分子生成再到分子逆合成规划）。在这些应用中，分子通常被表示成图的形式，以原子为节点，共价键为边。尽管这种表示方法很有效，但实际上最好的分子表示应该是三维结构（也就是“构象”），在三维结构中，每一个分子都有三维坐标与之对应。这种三维结构也更具内在性且信息含量高，包含了许多化学和生物学特性，如化学传感和与蛋白质的相互作用。

然而，通过实验的方式去确定分子的三维结构是很耗时费材的。有效预测有效构象和低能构象一直是计算化学中一个非常重要和活跃的课题。传统的计算化学方法是用基于经验的分子力场或量子学模拟计算自由能，再使用马尔可夫蒙特卡洛或分子动力学模拟去搜寻能量最低的构象。最近，更多的研究者利用机器学习方法去进行分子构象生成，常用的是二阶段的方法：1、基于原子图预测每一个原子对间的距离；2、依据预测的距离去生成分子构象。这种基于距离几何的方法有效地考虑了分子构象的旋转和平移不变性，因此获得了不错的性能。然而，两阶段的方法仍然有很大的局限性，它们分别预测距离和构象：预测的距离可能无法适当地保留保持真实原子的邻接关系。在第二阶段的距离几何程序可能会显著放大这些距离中的误差，从而产生异常的三维样本。

在本文中，我们提出了一种端到端基于双重优化的分子构象生成框架---ConfVAE。为了开发出具有旋转平移不变性的构象模型，我们仍然把原子对间距离作为中间变量。然而，我们没有学习通过最小化距离空间中的误差来预测距离，而是将整个问题表述为双重优化，同时优化了构象生成的距离预测问题和距离几何问题。整个框架建立在条件变分自编码器（CVAE）的基础上，在该框架中，分子图首先被编码到隐空间，并基于隐变量和分子图生成构象。在训练过程中，我们对距离预测模型中的一组距离进行迭代采样，通过最小化内部目标（由距离几何问题定义）生成三维结构，然后通过优化外部目标（即直接定义在构象上的可能性）来更新距离预测模型。

2 背景

2.1 体内毒物基因组学（TGx）数据集

每一个分子被表示为具有特征的图 G=，其中V代表原子的节点集，每一个点v都有对应的原子类别，E代表共价键的边的集合，每一条边 e_uv 都有对应的化学键类别。在本文中，我们还添加了辅助边以增强模型表现。对于原子几何 R，每一个点 v 还会被赋予三维坐标 r_v，d_uv 表示点 v 和点 u 之间的欧氏距离。

2.2 问题定义

分子构象生成问题是一个条件生成过程，其目标是在给定分子图的情况下对分子构象 R 的条件分布进行建模。

2.3 双重优化

双重优化是一种优化问题，其中（外部）目标函数中涉及的一组变量通过求解另一个（内部）优化问题获得。一般地，给定外部优化函数F及其参数 θ，内部优化函数 H 及其参数 w，双重优化可以表述为：

双重优化已经在超参优化、对抗学习、多任务学习等任务中显示出有效性。

通常，公式（1）是难解的，因为很难得到内部参数 w 的闭式解。通用的方法是用近似求解去替获得内部优化函数 H 的最小值，即使用T次随机梯度下降去迭代更新内部参数 w。一般来说，内外部参数是实值的，内外部函数是光滑的，所以可得到外部函数的超梯度：

3 隐式距离几何

在3.1部分，我们从宏观层面描述了我们的双重优化问题；在3.2部分，我们展示了模型的架构及训练的目标函数；在3.3部分，我们展示如何通过超梯度下降的方法去优化模型；3.4 部分展示了采样过程。

3.1 概述

因为一个分子可以有多个稳定构象，我们使用条件变分自编码器 CVAE，以分子图 G 为条件，建模了分子构象的分布 R，即p(R|G)。在 CVAE中，隐变量 z 被引入来建模分子构象生成的不确定性。已知关于隐变量的先验分布p(z|G)，编码器 q(z|R,G) 和解码器 p(R|z,G)，则可以表示 CVAE 的证据下界 L_ELBO：

其中，假设 p(z|G) 和 q(z|R,G) 符合相互独立的高斯分布，高斯分布的均值和方差可以通过作用在分子图上的图神经网络获得。CVAE就是要最大化证据下界 L_ELBO 来实现对目标分布 P(R|G) 的逼近。而在公式（3）中，L_ELBO 等于第一项重建损失 L_recon + 第二项先验损失 L_prior。在本文中，解码器为 D_θ(z,G)，即将隐变量z和分子图作为输入，输出所有原子对的距离。在下文我们会详细确立双重优化的问题：

内部优化。直接生成构象的三维坐标会严重受到随机旋转和平移的影响，例如：同一个构象可以有无数种平移旋转方式，导致其三维坐标不唯一。但是无论如何平移旋转，原子对间的距离是确定的。因此，ConfVAE要让解码器D_θ(z,G) 输出的是原子对间距离，并以此作为中间变量。为了生成一个分子构象 R，首先生成原子对距离 d，再对距离 d 进行后处理得到三维构象 R，此过程等价于求解距离几何优化问题：

此即为内部优化函数。

外部优化。最终，我们感兴趣的是直接最小化 3D 结构上的泛化误差，以使生成的构象在旋转和平移之前与真实构象一致。匹配后根均方差（RMSD）是一种广泛使用的度量标准。为了计算这个指标，首先需要将生成的构象通过匹配函数 A(·,·)，尽可能平移旋转到跟真实构象接近；再计算生成的构象和真实构象的 RMSD ：

因此，重建损失 L_recon，即外部优化函数可以写为：

双重优化。现在已经分别得到了公式（4）和公式（6）作为双重优化的内部、外部优化函数。在问题的构建中，外部优化函数旨在对真实的条件分布 p(R|G) 进行建模，而内部优化函数是解决了给定距离生成构象的问题。对公式（6）的隐变量 z 求期望，就可以得到公式（3）中重建损失 L_recon 的表示：

当下得到的双重优化还是很难解，因为：1）求解三维构象的内部优化问题不能得到闭式解；2）在隐变量空间中对隐变量求期望是难做到的。因此，在实践中，我们使用变分推断和重参数技巧对输出进行估计。我们将在以下部分详细阐述如何解决这些问题。

3.2 生成模型

信息传递神经网络（MPNN）是图神经网络的变种，此框架的所有模块（编码器，先验，解码器）都由 MPNN 实现。MPNN 直接作用于图G，对图具有同构不变性。在每个信息传递层中，通过聚合来自相邻节点的信息来更新原子编码。

对于编码器 q(z|R,G) 和先验 p(z|G)，我们使用相同的MPNN 框架。由于双重优化内存占用很大，我们使用基于常微分方程（ODE）的连续标准化流模型（CNF）去实现解码器 p(R|z,G)，这样解码器就只占用恒定的内存。

解码器结构。上文说到，解码器是二级级联结构：一个距离预测器 D_θ(z,G)，给定隐变量z和分子图 G 预测原子对间距离；一个可微的距离几何方程，从原子对间距离恢复三维构象 R。距离预测器 D_θ(z,G) 基于带条件的 CNF 实现：

因此，关于距离的条件对数似然可以表示为：

然后可以使用  ODE 求解器来估计参数上的梯度，以进行优化。在实践中，L_aux 可以作为一个辅助损失，根据距离定义，以监督训练。总之，总的损失函数可以解释为三个部分的总和：

整个模型的框架如图1所示：

图1 ConfVAE的整体框架

3.3 基于超梯度下降的端到端训练

以下阐述如何优化公式（7）（8）定义的双重优化问题。公式（8）所定义的内部优化，本质上是一个依据成对距离推断三维坐标的经典距离几何问题。在本文中，我们选取梯度下降，通过迭代的方式求解这个距离几何问题：

经过足够多次的迭代，依据给定的成对距离可以收敛到一个比较合适的空间构象 R。

现在我们考虑如何计算超梯度，即公式（7）目标函数的梯度，用来训练模型。假设三维构象R经过T次迭代之后收敛。则依据定义可以写出超梯度的公式：

其中，对于公式（13）第二行对于R的求导，由公式（12）迭代而成。相邻的一步迭代形如下式：

又由于 R 始于成对距离 d，因此可以计算 R 对 d 的导数：

算法的前向传播，反向传播的计算图如图2所示：

图2 算法的前向计算、反向计算示意图

训练过程的算法如图3所示：

图3 ConfVAE的训练算法

3.4 采样

给定一个分子图 G，需要生成对应的三维构象 R。我们先从先验分布 p(z|G) 中将隐变量 z 采样出来。然后我们从高斯分布中初始化距离，利用基于 ODE 的 CNF 得到原子对距离 d。最后我们去求解公式（4）的优化问题从而得到三维构象 R。

4 实验

本文的实验分为构象生成和距离分布建模，并选取当前最先进的基准模型进行比对。

在构象生成实验中，选取 GEOM-Drugs 数据学习；使用 RMSD 评判构象的相似度，使用覆盖分数Cov和匹配分数 MAT 分别评判生成构象的多样性和质量。构象生成的数据比对如表1所示，生成的部分构象如图4所示。

表1 构象生成的比对数据

图4 构象生成的可视化比较

在距离分布建模中，选取 ISO17 数据集进行学习，使用最大均值差异 MMD 作为评判指标，其值越小表明模型对距离的建模越精确。可以看出，ConfVAE 对距离建模的精确性是很高的。

表2 距离分布建模的数据比较

5 结论

在本文中，我们提出了 ConfVAE，一个通过双重优化生成分子构象的端到端框架。由于基于双重优化的端到端训练，我们的生成模型可以显著减小之前两阶段模型的误差，同时保持旋转和平移不变性。实验结果表明，在几个标准基准上，我们的方法优于所有最先进的模型。未来的工作包括将我们的双重优化框架与其他类型的生成模型相结合，并将我们的方法扩展到其他具有挑战性的结构，如蛋白质。

参考资料

Xu M, Wang W, Luo S, et al. An end-to-end framework for molecular conformation generation via bilevel programming[C]//International Conference on Machine Learning. PMLR, 2021: 11537-11547.

--------- End ---------

感兴趣的读者，可以添加小邦微信（zhiyaobang2020）加入读者实名讨论微信群。添加时请主动注明姓名-企业-职位/岗位 或  姓名-学校-职务/研究方向。

- 历史文章推荐 -

【分子生成】